# server节点需要产生一个张量作为输入
# 还需要转发数据包
import torch
import sys
import socket
import json
sys.path.append("..") 
import tensor2buff
from tensor2buff import HEADER_SIZE
from threading import Thread

class TCPServer(Thread):
    def __init__(self, port=8080):
        Thread.__init__(self)
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.bind(("0.0.0.0", port))
        self.socket.listen(128)
        self.clients = dict()
        
    def run(self):
        while True:
            client_socket , client_addr = self.socket.accept()
            print(client_addr, " accepted")
            self.clients[client_addr] = client_socket
            Thread(target=self.__client_task, args=(client_addr,)).start()
    
    # 每个client都起一个client_task线程
    def __client_task(self, client_addr):
        client_socket = self.clients[client_addr]
        def recvNbytes(N):
            cnt = N
            data = bytes()
            while cnt:
                buf = client_socket.recv(1024 if cnt > 1024 else cnt)
                if buf == b'': # 远程可能断线了 调用recvNbytes处注意检查len(data)
                    break
                data += buf
                cnt -= len(buf)
            return data
         
        while True:
            # TODO: read header
            header = recvNbytes(HEADER_SIZE)
            if len(header) != HEADER_SIZE:
                break
            info = json.loads(header.decode("utf-8"))
            # TODO: read pload
            zbuf = recvNbytes(info['length'])
            t = tensor2buff.unpack(info, zbuf)
            print(t)
            
        self.clients.pop(client_addr)
        print(f"{client_addr} disconnected")
    

if __name__ == "__main__":
    s = TCPServer()
    s.start()